/* uc65 - Micro Code for 6502/NES
 * Copyright (C) 2013  Norman B. Lancaster
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.gmail.qbradq.uc65;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.EnumSet;
import java.util.Set;

/**
 * The class responsible for generating output assembly code.
 */
public class CodeGenerator {
	private static final Set<Operation> readOps = EnumSet.of(
		Operation.ADC,
		Operation.AND,
		Operation.ASL,
		Operation.EOR,
		Operation.CMP,
		Operation.LDA,
		Operation.LDX,
		Operation.LDY,
		Operation.LSR,
		Operation.ORA,
		Operation.ROL,
		Operation.ROR,
		Operation.SBC
	);
	private static final Set<Operation> writeOps = EnumSet.of(
		Operation.STA
	);
	
	/**
	 * @param op The opcode to check.
	 * @return True if the opcode is a read operation, false otherwise.
	 */
	public static boolean isReadOp(Operation op) {
		return readOps.contains(op);
	}
	
	/**
	 * @param op The opcode to check.
	 * @return True if the opcode is a write operation, false otherwise.
	 */
	public static boolean isWriteOp(Operation op) {
		return writeOps.contains(op);
	}
	
	private BufferedWriter out;
	private Uc65 compiler;
	private File outFile;
	private boolean haveReportedWriteError = false;
	private boolean inRam = false;
	private boolean inRom = false;
	private int ramBank = -1;
	private int romBank = -1;
	private int nextLabel = 0;
	private Optimizer optimizer;
	
	/**
	 * Constructs a new CodeGenerator for a given source file.
	 * @param path The path of the source file.
	 * @param compiler The compiler to report errors to.
	 */
	public CodeGenerator(File source, Uc65 compiler) {
		optimizer = new Optimizer();
		this.compiler = compiler;
		
		outFile = new File(Uc65.outPath);
		try {
			out = new BufferedWriter(new FileWriter(outFile));
			out.write(String.format("; Assembly generated by uc65\n; File: %s\n", source.getPath()));
			if(Uc65.emitDebugInfo) {
				out.write(String.format(".DBG FILE, \"%s\", %d, %d\n", source.getPath(), source.length(), (int)(source.lastModified() / 1000L)));
			}
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * @param name The name to transform.
	 * @return A label-safe string for the given name.
	 */
	public String labelNameFrom(String name) {
		return name.replaceAll("[^a-zA-Z0-9_]", "_");
	}
	
	/**
	 * Closes the CodeGenerator and flushes output to disk.
	 */
	public void close() {
		try {
			out.flush();
			out.close();
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * @return An assembler label unique to this translation unit.
	 */
	public String label() {
		return String.format("_L%06X", nextLabel++);
	}
	
	private void reportWriteError() {
		if(!haveReportedWriteError) {
			compiler.error(String.format("Unable to write to output file %s", outFile.getPath()), null);
			haveReportedWriteError = true;
		}
	}
	
	private void switchRomBank(int bank) throws IOException {
		if(!inRom || bank != romBank) {
			inRom = true;
			inRam = false;
			romBank = bank;
			out.write(String.format(".SEGMENT \"ROM%d\"\n", bank));
		}
	}
	
	private void switchRamBank(int bank) throws IOException {
		if(!inRam || bank != ramBank) {
			inRom = false;
			inRam = true;
			ramBank = bank;
			out.write(String.format(".SEGMENT \"BSS%d\"\n", bank));
		}
	}
	
	/**
	 * Emits a label to the assembly output.
	 * @param name The label identifier.
	 */
	public void emitLabel(int romBank, String name) {
		try {
			switchRomBank(romBank);
			out.write(String.format("\t%s:\n", name));
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * Emits a variable label.
	 * @param line The source line that defined this variable.
	 * @param name The name of the variable.
	 * @param length The length of the variable, may be zero.
	 * @param extern Flag to say if the label should be exported.
	 * @param ramBank The ram bank to place this variable in.
	 */
	public void emitVariable(SourceLine line, String name, int length, boolean extern, int ramBank) {
		try {
			switchRamBank(ramBank);
			out.write(String.format("\t%s: .RES %d\n", name, length));
			if(extern) {
				out.write(String.format("\t.EXPORT %s\n", name));
			}
			if(Uc65.emitDebugInfo) {
				out.write(String.format(".DBG SYM, \"%s\", \"90\", %s, \"%s\"\n", name, extern ? "EXTERN" : "STATIC", name));
			}
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * Enters an assembly procedure block.
	 * @param line The source line the procedure was defined on.
	 * @param name The name of the procedure.
	 * @param extern Flag to say if the label should be exported.
	 * @param romBank The rom bank to place this procedure.
	 */
	public void enterProcedure(SourceLine line, String name, boolean extern, int romBank) {
		try {
			switchRomBank(romBank);
			emitSourceLine(line);
			out.write(String.format(".PROC %s\n", name));
			if(extern) {
				out.write(String.format(".EXPORT %s\n", name));
			}
			if(Uc65.emitDebugInfo) {
				out.write(String.format(".DBG FUNC, \"%s\", \"90\", %s, \"%s\"\n", name, extern ? "EXTERN" : "STATIC", name));
			}
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * Exits an assembly procedure block.
	 */
	public void exitProcedure() {
		try {
			out.write(".ENDPROC\n");
		} catch(IOException e) {
			reportWriteError();
		}		
	}
	
	/**
	 * Emits a source code line as a comment.
	 * @param line The source code line to emit.
	 */
	public void emitSourceLine(SourceLine line) {
		try {
			if(Uc65.emitSourceLines) {
				out.write(String.format(";%% %s\n", line.getText()));
			}
			if(Uc65.emitDebugInfo) {
				out.write(String.format(".DBG LINE, \"%s\", %d\n", line.getPath(), line.getLine()));
			}
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * Emits a source line to the output file verbatim.
	 * @param line The line to emit.
	 */
	public void emitVerbatimLine(SourceLine line) {
		try {
			emitSourceLine(line);
			out.write(String.format("\t%s\n", line.getText()));
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * Emits an "assign-style" label to a given address.
	 * @param name The name of the label.
	 * @param value The address to assign the label to.
	 * @param export Flag to see if this label should be exported.
	 */
	public void emitAssign(String name, int value, boolean export) {
		try {
			out.write(String.format("%s := $%04X\n", name, value));
			if(export) {
				out.write(String.format(".export %s\n", name));
			}
		} catch(IOException e) {
			reportWriteError();
		}
	}
	
	/**
	 * Called when a new, non-deterministic execution block begins.
	 */
	public void newExecutionBlock() {
		optimizer.reset();
	}
	
	/**
	 * Emits a single operation.
	 * @param op The opcode to emit.
	 * @param mode The addressing mode to use.
	 * @param romBank The ROM bank in which to place code.
	 * @param name An assembly label to use for addressed modes.
	 * @param value A byte value for immediate addressing.
	 * @param vol A flag to indicate if the instruction is volatile.
	 */
	public void emitOpCode(Operation op, AddressMode mode, int romBank, String name, byte value, boolean vol) {
		// If the optimizer tells us this instruction is redundant and it's not
		// a volatile instruction, don't emit it.
		if(!optimizer.instruction(op, mode, name, value) && !vol) {
			return;
		}
		String code = op.toString().toLowerCase();
		try {
			switchRomBank(romBank);
			switch(mode) {
				case IMPLICIT:
					out.write(String.format("\t%s\n", code));
					break;
				case IMMEDIATE:
					out.write(String.format("\t%s\t#$%02X\n", code, value));
					break;
				case ADDRESS:
					out.write(String.format("\t%s\t%s\n", code, name));
					break;
				case ADDRESSX:
					out.write(String.format("\t%s\t%s,x\n", code, name));
					break;
				case ADDRESSY:
					out.write(String.format("\t%s\t%s,y\n", code, name));
					break;
				case INDIRECT:
					out.write(String.format("\t%s\t(%s)\n", code, name));
					break;
				case INDIRECTX:
					out.write(String.format("\t%s\t(%s,x),x\n", code, name));
					break;
				case INDIRECTY:
					out.write(String.format("\t%s\t(%s),y,y\n", code, name));
					break;
			}
		} catch(IOException e) {
			reportWriteError();
		}
	}
}
